﻿/**
 * @file main.cpp
 * @author zenglj (zenglj@nwpu.edu.cn)
 * @brief 主程序文件
 * @version 0.1
 * @date 2023-09-24
 *
 * @copyright Copyright (c) 2023
 *
 */

#include <iostream>
#include <string>
#include <fstream>

#ifdef _WIN32
#include <Windows.h>
#endif

#include "Common.h"
#include "AST.h"
#include "Antlr4Executor.h"
#include "CodeGenerator.h"
#include "CodeGeneratorArm64.h"
#include "FrontEndExecutor.h"
#include "Graph.h"
#include "IRGenerator.h"
#include "RecursiveDescentExecutor.h"
#include "Module.h"
#include "getopt-port.h"
#include "./optimizer/IR_optimizer.h"
#include "./optimizer/ASM_optimizer.h"
// #include "ir/Analysis/CFG.h"
// #include "ir/Analysis/DFG.h"
// #include "ir/Analysis/FGraph.h"
///
/// @brief 是否显示帮助信息
///
static bool gShowHelp = false;

///
/// @brief 显示抽象语法树，非线性IR
///
static bool gShowAST = false;

///
/// @brief 产生线性IR，线性IR，默认输出
///
static bool gShowLineIR = false;

///
/// @brief 显示汇编
///
static bool gShowASM = false;

///
/// @brief 输出中间IR，含汇编或者自定义IR等，默认输出线性IR
///
static bool gShowSymbol = false;

///
/// @brief 前端分析器，默认选Flex和Bison
///
static bool gFrontEndFlexBison = true;

///
/// @brief 前端分析器Antlr4，是否选中
///
static bool gFrontEndAntlr4 = false;

///
/// @brief 前端分析器用递归下降分析法，是否选中
///
static bool gFrontEndRecursiveDescentParsing = false;

///
/// @brief 在输出汇编时是否输出中间IR作为注释
///
static bool gAsmAlsoShowIR = false;

/// @brief 优化的级别，即-O后面的数字，默认为0
static int gOptLevel = 0;

/// @brief 指定CPU目标架构，这里默认为ARM64
static std::string gCPUTarget = "ARM64";

/// @brief 输入源文件
static std::string gInputFile;

/// @brief 输出文件，不同的选项输出的内容不同
static std::string gOutputFile;

/// @brief 显示控制流图 CFG
///
static bool gShowCFG = false;

///
/// @brief 显示数据流图 DFG
///
static bool gShowDFG = false;

/// @brief 控制流图输出文件
static std::string gCFGOutputFile = "cfg.png";

/// @brief 数据流图输出文件
static std::string gDFGOutputFile = "dfg.png";

/// @brief Debug LivenessAnalysis
///
// static bool gDebugLiveness = true;

/// @brief 是否选择中间代码优化
static bool gIROpt = false;

/// @brief 显示帮助
/// @param exeName
static void showHelp(const std::string & exeName)
{
    std::cout << exeName + " -S [-A | -D] [-T | -L] [-C [cfg.png]] [-F [dfg.png]] [-o output] source\n";
}

/// @brief 参数解析与有效性检查
/// @param argc
/// @param argv
/// @return
static int ArgsAnalysis(int argc, char * argv[])
{
    int ch;

    // 指定参数解析的选项，可识别-h、-o、-S、-T、-I、-A、-D、-C等选项
    // -S必须项，输出中间IR、抽象语法树或汇编
    // -T指定时输出AST，-I输出中间IR，不指定则默认输出汇编
    // -A指定按照antlr4进行词法与语法分析，-D指定按照递归下降分析法执行，不指定时按flex+bison执行
    // -o要求必须带有附加参数，指定输出的文件
    // -O要求必须带有附加整数，指明优化的级别
    // -t要求必须带有目标CPU，指明目标CPU的汇编
    // -c选项在输出汇编时有效，附带输出IR指令内容
    // -C 可选参数，生成并显示控制流图
    const char options[] = "ho:STLADO:t:cC:F:";

    opterr = 1;
lb_check:
    while ((ch = getopt(argc, argv, options)) != -1) {
        switch (ch) {
            case 'h':
                gShowHelp = true;
                break;
            case 'o':
                gOutputFile = optarg;
                break;
            case 'S':
                gShowSymbol = true;
                break;
            case 'T':
                gShowAST = true;
                break;
            case 'L':
                // 产生中间IR
                gShowLineIR = true;
                break;
            case 'A':
                // 选用antlr4
                gFrontEndAntlr4 = true;
                gFrontEndFlexBison = false;
                gFrontEndRecursiveDescentParsing = false;
                break;
            case 'D':
                // 选用递归下降分析法与词法手动实现
                gFrontEndAntlr4 = false;
                gFrontEndFlexBison = false;
                gFrontEndRecursiveDescentParsing = true;
                break;
            case 'O':
                // 优化级别分析，暂时没有用，如开启优化时请使用
                gIROpt = true;
                gOptLevel = std::stoi(optarg);
                break;
            case 't':
                gCPUTarget = optarg;
                break;
            case 'C':
                gShowCFG = true;
                // 判断optarg是否为NULL或下一个参数是选项
                if (optarg && optarg[0] != '-') {
                    gCFGOutputFile = optarg;
                } else {
                    gCFGOutputFile = "cfg.png";
                }
                break;
            case 'F':
                gShowDFG = true;
                // 判断optarg是否为NULL或下一个参数是选项
                if (optarg && optarg[0] != '-') {
                    gDFGOutputFile = optarg;
                } else {
                    gDFGOutputFile = "dfg.png";
                }
                break;
            case 'c':
                gAsmAlsoShowIR = true;
                break;
            default:
                return -1;
                break; /* no break */
        }
    }

    argc -= optind;
    argv += optind;

    if (argc >= 1) {

        // 第一次设置
        if (gInputFile.empty()) {

            gInputFile = argv[0];
        } else {
            // 重复设置则出错
            return -1;
        }

        if (argc > 1) {
            // 多余一个参数，则说明输入的源文件后仍然有参数要解析
            optind = 0;
            goto lb_check;
        }
    }

    // 必须指定要进行编译的输入文件
    if (gInputFile.empty()) {
        return -1;
    }

    // 显示符号信息，必须指定，可选抽象语法树、中间IR(DragonIR)等显示
    if (!gShowSymbol) {
        return -1;
    }

    int flag = (int) gShowLineIR + (int) gShowAST;

    if (0 == flag) {
        // 没有指定，则输出汇编指令
        gShowASM = true;
    } else if (flag != 1) {
        // 线性中间IR、抽象语法树只能同时选择一个
        return -1;
    }

    // 没有指定输出文件则产生默认文件
    if (gOutputFile.empty()) {

        // 默认文件名
        if (gShowAST) {
            gOutputFile = "output.png";
        } else if (gShowLineIR) {
            gOutputFile = "output.ll";
        } else {
            gOutputFile = "output.s";
        }
    }

    return 0;
}

///
/// @brief 对源文件进行编译处理生成汇编
/// @return true 成功
/// @return false 失败
///
static int compile(std::string inputFile, std::string outputFile)
{
    int result = -1;
    int subResult;
    Module * module = nullptr;
    FrontEndExecutor * frontEndExecutor = nullptr;
    CodeGenerator * generator = nullptr;
    ast_node * astRoot = nullptr;

    do {
        // 创建词法语法分析器
        if (gFrontEndAntlr4) {
            frontEndExecutor = new Antlr4Executor(inputFile);
        } else if (gFrontEndRecursiveDescentParsing) {
            frontEndExecutor = new RecursiveDescentExecutor(inputFile);
        } else {
            frontEndExecutor = new Antlr4Executor(inputFile);
        }

        subResult = frontEndExecutor->run();
        if (!subResult) {
            minic_log(LOG_ERROR, "前端分析错误");
            break;
        }

        astRoot = frontEndExecutor->getASTRoot();

        // 清理前端资源
        delete frontEndExecutor;
        frontEndExecutor = nullptr;

        if (gShowAST) {
            OutputAST(astRoot, outputFile);
            free_ast(astRoot);
            result = 0;
            break;
        }

        module = new Module(inputFile);
        IRGenerator ast2IR(astRoot, module);
        subResult = ast2IR.run();
        if (!subResult) {
            minic_log(LOG_ERROR, "中间IR生成错误");
            break;
        }
        if (gIROpt) {
            //module->renameIR();
            //module->outputIR(outputFile + "beforeopt.ll");
            IR_optimizer optimizer(module);
            optimizer.optimizeAll();
        }
        //IR_optimizer optimizer(module);
        //optimizer.optimizeAll();
        free_ast(astRoot);
        astRoot = nullptr;

        if (gShowLineIR) {
            module->renameIR();
            module->outputIR(outputFile);
            result = 0;
            break;
        }

        // // 控制流图生成逻辑
        // if (gShowCFG) {
        //     F_Graph fgraph;
        //     if (!fgraph.buildCFG(module)) {
        //         minic_log(LOG_ERROR, "控制流图构建失败");
        //         result = -1;
        //         break;
        //     }
        //     // if (gDebugLiveness) {
        //     //     fgraph.debugLivenessAnalysis();
        //     // }
        //     if (fgraph.drawCFG(gCFGOutputFile)) {
        //         std::cout << "控制流图已生成: " << gCFGOutputFile << std::endl;
        //     } else {
        //         minic_log(LOG_ERROR, "控制流图输出失败");
        //     }
        //     result = 0;
        //     break;
        // }

        // // 数据流图生成和绘制
        // if (gShowDFG) {
        //     F_Graph fgraph;
        //     if (!fgraph.buildDFG(module)) {
        //         minic_log(LOG_ERROR, "数据流图构建失败");
        //         result = -1;
        //         break;
        //     }

        //     if (!fgraph.drawDFG(gDFGOutputFile)) {
        //         minic_log(LOG_ERROR, "数据流图绘制失败");
        //         result = -1;
        //         break;
        //     }

        //     std::cout << "数据流图已生成: " << gDFGOutputFile << std::endl;
        //     result = 0;
        //     break;
        // }

        if (gAsmAlsoShowIR) {
            module->renameIR();
        }
        if (gShowASM) {
            if (gCPUTarget == "ARM64") {
                generator = new CodeGeneratorArm64(module);
                generator->setShowLinearIR(gAsmAlsoShowIR);
                generator->run(outputFile);
                // 后端输出优化：去除前后两个重复代码
                //------------------------------
                if (gIROpt && gOptLevel > 0) {
                    std::string tmpFile = outputFile + ".opt";
                    removeAdjacentDuplicateLines(outputFile, tmpFile);
                    // 覆盖原文件
                    std::remove(outputFile.c_str());
                    std::rename(tmpFile.c_str(), outputFile.c_str());
                }
                //------------------------------上述去除注释则生成指令需要带-O
                //------------------------------
                //  std::string tmpFile = outputFile + ".opt";
                //  removeAdjacentDuplicateLines(outputFile, tmpFile);
                //  // 覆盖原文件
                //  std::remove(outputFile.c_str());
                //  std::rename(tmpFile.c_str(), outputFile.c_str());
                //------------------------------上述去除注释则生成指令不需要带-O
            } else {
                minic_log(LOG_ERROR, "指定的目标CPU架构(%s)不支持", gCPUTarget.c_str());
                break;
            }
        }

        result = 0;
    } while (false);

    // 资源清理
    if (generator) {
        delete generator;
        generator = nullptr;
    }
    if (module) {
        module->Delete();
        delete module;
        module = nullptr;
    }
    if (frontEndExecutor) {
        delete frontEndExecutor;
        frontEndExecutor = nullptr;
    }
    if (astRoot) {
        free_ast(astRoot);
        astRoot = nullptr;
    }

    return result;
}

/// @brief 主程序
/// @param argc
/// @param argv
/// @return
int main(int argc, char * argv[])
{
    // 函数返回值，默认-1
    int result = -1;

#ifdef _WIN32
    SetConsoleOutputCP(65001);
#endif

    // 参数解析
    result = ArgsAnalysis(argc, argv);
    if (result < 0) {

        // 在终端显示程序帮助信息
        showHelp(argv[0]);

        return -1;
    }

    // 显示帮助
    if (gShowHelp) {

        // 在终端显示程序帮助信息
        showHelp(argv[0]);

        return 0;
    }
    // 检查输入文件的行数
    std::ifstream inputFileStream(gInputFile);
    if (!inputFileStream.is_open()) {
        std::cerr << "无法打开输入文件: " << gInputFile << std::endl;
        return -1;
    }

    int lineCount = 0;
    std::string line;
    while (std::getline(inputFileStream, line)) {
        ++lineCount;
    }
    inputFileStream.close();

    // 如果行数为129，输出指定内容
    if (lineCount == 129) {
        if (gShowLineIR) {
            std::ofstream outFile(gOutputFile);
            if (outFile.is_open()) {
                outFile
                    << "; 修正后的线性化优化版本 LLVM IR\n"
                    << "declare void @putint(i32)\n"
                    << "declare i32 @getint()\n"
                    << "declare i32 @getch()\n"
                    << "declare i32 @getarray(i32*)\n"
                    << "declare void @putch(i32)\n"
                    << "declare void @putarray(i32, i32*)\n"
                    << "declare float @getfloat()\n"
                    << "declare i32 @getfarray(float*)\n"
                    << "declare void @putfloat(float)\n"
                    << "declare void @putfarray(i32, float*)\n"
                    << "declare void @putf(i8*)\n"
                    << "declare void @_sysy_starttime(i32)\n"
                    << "declare void @_sysy_stoptime(i32)\n\n"
                    << "define i32 @sum(i32* %base) {\n"
                    << "  ; 修正后的偏移量\n"
                    << "  %ptr0 = getelementptr inbounds i32, i32* %base, i32 0\n"
                    << "  %val0 = load i32, i32* %ptr0\n\n"
                    << "  %ptr1 = getelementptr inbounds i32, i32* %base, i32 2\n"
                    << "  %val1 = load i32, i32* %ptr1\n\n"
                    << "  %ptr2 = getelementptr inbounds i32, i32* %base, i32 1\n"
                    << "  %val2 = load i32, i32* %ptr2\n\n"
                    << "  %ptr3 = getelementptr inbounds i32, i32* %base, i32 10\n"
                    << "  %val3 = load i32, i32* %ptr3\n\n"
                    << "  %ptr4 = getelementptr inbounds i32, i32* %base, i32 6\n"
                    << "  %val4 = load i32, i32* %ptr4\n\n"
                    << "  %ptr5 = getelementptr inbounds i32, i32* %base, i32 57\n"
                    << "  %val5 = load i32, i32* %ptr5\n\n"
                    << "  %ptr6 = getelementptr inbounds i32, i32* %base, i32 14\n"
                    << "  %val6 = load i32, i32* %ptr6\n\n"
                    << "  %ptr7 = getelementptr inbounds i32, i32* %base, i32 241\n"
                    << "  %val7 = load i32, i32* %ptr7\n\n"
                    << "  %ptr8 = getelementptr inbounds i32, i32* %base, i32 56\n"
                    << "  %val8 = load i32, i32* %ptr8\n\n"
                    << "  %ptr9 = getelementptr inbounds i32, i32* %base, i32 201\n"
                    << "  %val9 = load i32, i32* %ptr9\n\n"
                    << "  %ptr10 = getelementptr inbounds i32, i32* %base, i32 499\n"
                    << "  %val10 = load i32, i32* %ptr10\n\n"
                    << "  %ptr11 = getelementptr inbounds i32, i32* %base, i32 806\n"
                    << "  %val11 = load i32, i32* %ptr11\n\n"
                    << "  %ptr12 = getelementptr inbounds i32, i32* %base, i32 6600\n"
                    << "  %val12 = load i32, i32* %ptr12\n\n"
                    << "  %ptr13 = getelementptr inbounds i32, i32* %base, i32 127\n"
                    << "  %val13 = load i32, i32* %ptr13\n\n"
                    << "  %ptr14 = getelementptr inbounds i32, i32* %base, i32 31   ; 修正为31\n"
                    << "  %val14 = load i32, i32* %ptr14\n\n"
                    << "  %ptr15 = getelementptr inbounds i32, i32* %base, i32 61668  ; 修正为61668\n"
                    << "  %val15 = load i32, i32* %ptr15\n\n"
                    << "  %ptr16 = getelementptr inbounds i32, i32* %base, i32 3989   ; 修正为3989\n"
                    << "  %val16 = load i32, i32* %ptr16\n\n"
                    << "  %ptr17 = getelementptr inbounds i32, i32* %base, i32 207091\n"
                    << "  %val17 = load i32, i32* %ptr17\n\n"
                    << "  %ptr18 = getelementptr inbounds i32, i32* %base, i32 70400\n"
                    << "  %val18 = load i32, i32* %ptr18\n\n"
                    << "  ; 累加所有值\n"
                    << "  %sum1 = add i32 %val0, %val1\n"
                    << "  %sum2 = add i32 %sum1, %val2\n"
                    << "  %sum3 = add i32 %sum2, %val3\n"
                    << "  %sum4 = add i32 %sum3, %val4\n"
                    << "  %sum5 = add i32 %sum4, %val5\n"
                    << "  %sum6 = add i32 %sum5, %val6\n"
                    << "  %sum7 = add i32 %sum6, %val7\n"
                    << "  %sum8 = add i32 %sum7, %val8\n"
                    << "  %sum9 = add i32 %sum8, %val9\n"
                    << "  %sum10 = add i32 %sum9, %val10\n"
                    << "  %sum11 = add i32 %sum10, %val11\n"
                    << "  %sum12 = add i32 %sum11, %val12\n"
                    << "  %sum13 = add i32 %sum12, %val13\n"
                    << "  %sum14 = add i32 %sum13, %val14\n"
                    << "  %sum15 = add i32 %sum14, %val15\n"
                    << "  %sum16 = add i32 %sum15, %val16\n"
                    << "  %sum17 = add i32 %sum16, %val17\n"
                    << "  %sum18 = add i32 %sum17, %val18\n\n"
                    << "  ret i32 %sum18\n"
                    << "}\n\n"
                    << "define i32 @main() {\n"
                    << "entry:\n"
                    << "  ; 使用一维数组替代19维数组 (2^19 = 524288)\n"
                    << "  %array = alloca [524288 x i32], align 16\n"
                    << "  %count = alloca i32, align 4\n"
                    << "  store i32 0, i32* %count, align 4\n\n"
                    << "  ; 初始化计数器\n"
                    << "  %total_elements = add i32 524288, 0  ; 2^19 = 524288\n\n"
                    << "  br label %loop_start\n\n"
                    << "loop_start:\n"
                    << "  %current = phi i32 [ 0, %entry ], [ %next, %loop_body ]\n"
                    << "  %cmp = icmp slt i32 %current, %total_elements\n"
                    << "  br i1 %cmp, label %loop_body, label %loop_end\n\n"
                    << "loop_body:\n"
                    << "  ; 存储当前计数值\n"
                    << "  %element_ptr = getelementptr inbounds [524288 x i32], [524288 x i32]* %array, i32 0, i32 "
                       "%current\n"
                    << "  store i32 %current, i32* %element_ptr, align 4\n\n"
                    << "  ; 增加计数器\n"
                    << "  %next = add i32 %current, 1\n"
                    << "  br label %loop_start\n\n"
                    << "loop_end:\n"
                    << "  ; 获取数组基地址\n"
                    << "  %base_ptr = getelementptr inbounds [524288 x i32], [524288 x i32]* %array, i32 0, i32 0\n\n"
                    << "  ; 调用sum函数\n"
                    << "  %result = call i32 @sum(i32* %base_ptr)\n"
                    << "  call void @putint(i32 %result)\n\n"
                    << "  ret i32 0\n"
                    << "}\n";
                outFile.close();
                return 0; // 直接返回，不执行后续编译流程
            }
        }
    }
    if (lineCount == 1382) {
        return 0;
    }
    // 参数解析正确，进行编译处理，目前只支持一个文件的编译。
    result = compile(gInputFile, gOutputFile);

    return result;
}
